###########################################################################
## Attributions of Fault and Support for Government Assistance for Workers
## Displaced by Technology
###########################################################################
## Seth Werfel
## Christopher Witko
## Tobias Heinrich
###########################################################################
## 2) Estimation
###########################################################################


## Load prepared data
#####################
load(file="output/Data_prepped.Rdata")
n_sim <- 1000


## Simple treatment effects; ordered probit
###########################################
## Storage
pred <- array(NA, dim=c(2, 3, n_sim, 7), dimnames=list(Y=c("Fault", "Support"),
                                                       Treat=c("Domestic", "Foreign", "Tech"),
                                                       Iter=NULL, K=1:7))
diff <- pred

## Estimation of ordered probits
mod_M <- bayespolr(WFault ~ 1 +  Treat + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
                     + Ideology_conservative + Ideology_liberal +
                     + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
                   data=data, weights=weight, method="probit")
mod_Y <- bayespolr(Support~ 1 + Treat  + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
                     + Ideology_conservative + Ideology_liberal +
                     + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
                   data=data, weights=weight, method="probit")


## Save coefficients
coefs <- array(NA, dim=c(2, length(c(coef(mod_M), mod_M$zeta)), n_sim),
               dimnames=list(Y=c("Fault", "Support"), Names=names(c(coef(mod_M), mod_M$zeta)),
                             Iter=NULL))
coefs["Fault", , ] <- t(rmvnorm(n_sim, c(coef(mod_M), mod_M$zeta), vcov(mod_M)))
coefs["Support", , ] <- t(rmvnorm(n_sim, c(coef(mod_Y), mod_Y$zeta), vcov(mod_Y)))


## Calculate predictions/ effects
for(j in 1:2)
{
  for(i in 1:3)
  {
    if(j == 1) theta <- rmvnorm(n_sim, c(coef(mod_M), mod_M$zeta), vcov(mod_M))
    if(j == 2) theta <- rmvnorm(n_sim, c(coef(mod_Y), mod_Y$zeta), vcov(mod_Y))
    
    beta <- theta[, 1:length(coef(mod_M))]
    zeta <- cbind(-Inf, theta[, (length(coef(mod_M))+1):ncol(theta)], Inf)
    J <- ncol(zeta) - 1
    
    tmp <- data
    X <- model.matrix(formula(mod_Y$terms), data=data)
    X[, c("TreatForeign", "TreatTech")] <- 0
    
    if(i == 2) X[, "TreatForeign"] <- 1
    if(i == 3) X[, "TreatTech"] <- 1
    X <- X[, -1]
    
    lp <- (beta) %*% t(X)
    
    for(k in 2:(J+1))
    {
      tmp <- pnorm(zeta[, k] - lp) - pnorm(zeta[, k-1] - lp)
      pred[j, i, , k-1] <- apply(tmp, 1, weighted.mean, data$weight)
    }
  }
  
  ## Subtract "Domestic" to get treatment effect
  for(i in 1:3)
  {
    diff[j, i, , ] <- pred[j, i, , ] - pred[j, 1, , ]
  }
}


## Summarize
main_pred <- adply(.data=pred, .margins=c(1, 2, 4),
                   .fun=function(x) data.frame(Mean=mean(x), Q025=quantile(x, 0.025),
                                               Q975=quantile(x, 0.975)))
main_diff <- adply(.data=diff, .margins=c(1, 2, 4),
                   .fun=function(x) data.frame(Mean=mean(x), Q025=quantile(x, 0.025),
                                               Q975=quantile(x, 0.975),
                                               PP=mean(x > 0)))
main_coefs <- adply(.data=coefs, .margins=c(1, 2),
                    .fun=function(x) data.frame(Mean=mean(x), Q025=quantile(x, 0.025),
                                                Q975=quantile(x, 0.975),
                                                PP=mean(x > 0)))



## Causal mediation
###################
## Storage
store_med <- array(NA, dim=c(n_sim, 2, 5, 2, 7),
                   dimnames=list(Iter=NULL, Treat=c("Foreign", "Tech"), 
                                 Comparison=c("Main", "Ideology_liberal", "Ideology_conservative", "FamInc_high",
                                              "FamInc_low"),
                                 Estimand=c("Direct", "Indirect"),
                                 Y=1:7))

## Estimate the necessary models
mod_M_simple <- bayespolr(WFault ~ 1 +  Treat + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
                            + Ideology_conservative + Ideology_liberal +
                            + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
                          data=data, weights=weight, method="probit")

mod_Y_simple <- bayespolr(Support~ 1 + Treat + WFault + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
                            + Ideology_conservative + Ideology_liberal +
                            + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
                          data=data, weights=weight, method="probit")

mod_Y_M_and_Ideology <- bayespolr(Support~ 1 + Treat+ WFault*Ideology_liberal+
                                    + WFault*Ideology_conservative + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
                                    + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
                                  data=data, weights=weight, method="probit")

mod_Y_M_and_FamInc <- bayespolr(Support~ 1 + Treat+ WFault*FamilyInc + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
                                  +  Ideology_conservative + Ideology_liberal +
                                  + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
                                data=data, weights=weight, method="probit")

# mod_Y_MT_and_Education <- bayespolr(Support~ 1 + Treat*Education_uni + WFault*Education_uni + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
#                                       +  Ideology_conservative + Ideology_liberal +
#                                       + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
#                                     data=data, weights=weight, method="probit")
# 
# mod_M_T_and_Education <- bayespolr(WFault ~ 1 +  Treat*Education_uni + Employed_full + Race_white + Race_black + Race_hispanic + Education_uni +
#                                      +  Ideology_conservative + Ideology_liberal +
#                                      + pid3 + Unemployed + FamilyInc + VoterReg_yes + Gender_male + NewsInt_high + Age,
#                                    data=data, weights=weight, method="probit")


## Store key coefficients
med_coefs <- rmvnorm(n_sim, c(coef(mod_Y_simple), mod_Y_simple$zeta), vcov(mod_Y_simple))


## Estimate mediation effects
for(i in 1:2)
{
  this_non_dom <- ifelse(i == 1, "Foreign", "Tech")
  
  ## Main analysis
  out_main <- mediate(model.m=mod_M_simple, model.y=mod_Y_simple,
                      sims=n_sim, treat="Treat", 
                      treat.value=this_non_dom, control.value="Domestic",
                      mediator="WFault")
  
  store_med[ , this_non_dom, "Main", "Direct", ] <- 0.5 * (out_main$z0.sims + out_main$z1.sims)
  store_med[ , this_non_dom, "Main", "Indirect", ] <- 0.5 * (out_main$d0.sims + out_main$d1.sims)
  
  
  ## Ideology analysis
  out_aux_ideology_liberal <- mediate(model.m=mod_M_simple, model.y=mod_Y_M_and_Ideology,
                                      sims=n_sim, treat="Treat", 
                                      treat.value=this_non_dom, control.value="Domestic",
                                      mediator="WFault", covariates=list(Ideology_conservative=0, Ideology_liberal=1))
  out_aux_ideology_conservative <- mediate(model.m=mod_M_simple, model.y=mod_Y_M_and_Ideology,
                                           sims=n_sim, treat="Treat", 
                                           treat.value=this_non_dom, control.value="Domestic",
                                           mediator="WFault", covariates=list(Ideology_conservative=1, Ideology_liberal=0))
  
  store_med[ , this_non_dom, "Ideology_liberal", "Direct", ] <- 0.5 * (out_aux_ideology_liberal$z0.sims +out_aux_ideology_liberal$z1.sims)
  store_med[ , this_non_dom, "Ideology_liberal", "Indirect", ] <- 0.5 * (out_aux_ideology_liberal$d0.sims +out_aux_ideology_liberal$d1.sims)
  store_med[ , this_non_dom, "Ideology_conservative", "Direct", ] <- 0.5 * (out_aux_ideology_conservative$z0.sims + out_aux_ideology_conservative$z1.sims)
  store_med[ , this_non_dom, "Ideology_conservative", "Indirect", ] <- 0.5 * (out_aux_ideology_conservative$d0.sims + out_aux_ideology_conservative$d1.sims)
  
  
  ## Family income analysis
  out_aux_faminc_high <- mediate(model.m=mod_M_simple, model.y=mod_Y_M_and_FamInc,
                                 sims=n_sim, treat="Treat", 
                                 treat.value=this_non_dom, control.value="Domestic",
                                 mediator="WFault", covariates=list(FamilyInc="High"))
  out_aux_faminc_low <- mediate(model.m=mod_M_simple, model.y=mod_Y_M_and_FamInc,
                                sims=n_sim, treat="Treat", 
                                treat.value=this_non_dom, control.value="Domestic",
                                mediator="WFault", covariates=list(FamilyInc="Low"))
  
  store_med[ , this_non_dom, "FamInc_high", "Direct", ] <- 0.5 * (out_aux_faminc_high$z0.sims + out_aux_faminc_high$z1.sims)
  store_med[ , this_non_dom, "FamInc_high", "Indirect", ] <- 0.5 * (out_aux_faminc_high$d0.sims + out_aux_faminc_high$d1.sims)
  store_med[ , this_non_dom, "FamInc_low", "Direct", ] <- 0.5 * (out_aux_faminc_low$z0.sims +  out_aux_faminc_low$z1.sims)
  store_med[ , this_non_dom, "FamInc_low", "Indirect", ] <- 0.5 * (out_aux_faminc_low$d0.sims + out_aux_faminc_low$d1.sims)
  
  
  # ## Education
  # out_aux_education_low <- mediate(model.m=mod_M_T_and_Education, model.y=mod_Y_MT_and_Education,
  #                                  sims=n_sim, treat="Treat", 
  #                                  treat.value=this_non_dom, control.value="Domestic",
  #                                  mediator="WFault", covariates=list(Education_uni=0))
  # out_aux_education_high <- mediate(model.m=mod_M_T_and_Education, model.y=mod_Y_MT_and_Education,
  #                                   sims=n_sim, treat="Treat", 
  #                                   treat.value=this_non_dom, control.value="Domestic",
  #                                   mediator="WFault", covariates=list(Education_uni=1))
  # 
  # store_med[ , this_non_dom, "Edu_high", "Direct", ] <- 0.5 * (out_aux_education_high$z0.sims + out_aux_education_high$z1.sims)
  # store_med[ , this_non_dom, "Edu_high", "Indirect", ] <- 0.5 * (out_aux_education_high$d0.sims + out_aux_education_high$d1.sims)
  # 
  # store_med[ , this_non_dom, "Edu_low", "Direct", ] <- 0.5 * (out_aux_education_low$z0.sims + out_aux_education_low$z1.sims)
  # store_med[ , this_non_dom, "Edu_low", "Indirect", ] <- 0.5 * (out_aux_education_low$d0.sims + out_aux_education_low$d1.sims)
}


## Summarize
store_med <- adply(.data=store_med, .margins=2:5,
                   .fun=function(x) data.frame(Mean=mean(x), Q025=quantile(x, 0.025),
                                               Q975=quantile(x, 0.975), PP=mean(x > 0)))

med_coefs <- adply(.data=med_coefs, .margins=2,
                   .fun=function(x) data.frame(Mean=mean(x), Q025=quantile(x, 0.025),
                                               Q975=quantile(x, 0.975), Y="Support_M",PP=mean(x > 0)))
colnames(med_coefs)[1] <- "Names"



## Save all to disk
####################
## Mediation results
save(store_med, file="output/Mediation_results.Rdata")

## Main effects/ predictions
save(main_diff, file="output/Main_difference.Rdata")
save(main_pred, file="output/Main_predictions.Rdata")

## Save coefficients
all_coefs <- rbind(med_coefs, main_coefs)
save(all_coefs, file="output/All_coefficients.Rdata")



## Garbage collection
#####################
rm(list=ls())

